from utils.common_import import *
from .model import PlatformUsers
from .model import PAI_DATASET, PAI_PROJECT, PAI_PROJECT_JOB

from config.cluster_config_new import cluster_client
gpu_device_name = cluster_client.cluster_config["device"]["gpu"]["resource_name"]

class OverView(Resource):
    # @jwt_required()
    @exception_response()
    def get(self):
        #获取集群所有任务
        nodes_list = cluster_client.kube_client.get_nodes().items
        total_nodes = len(nodes_list)
        logger.info("total_nodes is {}".format(total_nodes))

        total_users = PlatformUsers.query.count()
        logger.info("total_users is {}".format(total_users))

        # node_gpu_total_exp = "count(nvidiasmi_utilization_gpu)"
        # total_gpus_result = cluster_client.prometheus_client.query(node_gpu_total_exp)
        # logger.info("total_gpus is {}".format(total_gpus_result))
        # total_gpus = self._get_count_from_prometheus_result(total_gpus_result)

        total_gpus = 0
        for node in nodes_list:
            total_gpus += float(node.status.capacity.get(gpu_device_name, 0))

        total_datasets = PAI_DATASET.query.count()
        logger.info("total_datasets is {}".format(total_datasets))

        total_public_projects = PAI_PROJECT.query.filter_by(type=1).count()
        logger.info("total_public_projects is {}".format(total_public_projects))

        total_jobs = PAI_PROJECT_JOB.query.count()
        logger.info("total_jobs is {}".format(total_jobs))

        alert_query_Name = 'sum(ALERTS{alertstate="firing"})'
        total_alerts_result = cluster_client.prometheus_client.query(alert_query_Name)
        logger.info("total_alerts_result is {}".format(total_alerts_result))

        total_alerts = self._get_count_from_prometheus_result(total_alerts_result)

        overview_info = dict(
            total_nodes = total_nodes,
            total_users = total_users,
            total_gpus = total_gpus,
            total_datasets = total_datasets,
            total_public_projects = total_public_projects,
            total_jobs = total_jobs,
            total_alerts = total_alerts
        )

        return make_success_response(overview_info)

    def _get_count_from_prometheus_result(self, prometheus_result):
        if prometheus_result \
                and 'data' in prometheus_result \
                and 'result' in prometheus_result["data"]\
                and prometheus_result['data']['result']:
            result_count = prometheus_result['data']['result'][0]['values'][0][1]
        else:
            result_count = 0

        return result_count
